#!/usr/bin/env python3
"""
Consul服务管理脚本
用于查看和清理失效的服务实例

使用方法:
python consul_manager.py --list                          # 列出所有服务
python consul_manager.py --list-critical                 # 列出所有critical状态的服务  
python consul_manager.py --cleanup                       # 清理所有critical状态的服务
python consul_manager.py --cleanup-service user_api      # 清理指定服务的critical实例
python consul_manager.py --cleanup-duplicates            # 清理重复的服务实例（推荐）
python consul_manager.py --force-cleanup user_api        # 强制清理指定服务的所有实例（危险）
python consul_manager.py --test-discovery                # 测试服务发现功能
python consul_manager.py --show-failed-count             # 显示当前失效服务计数状态
python consul_manager.py --reset-failed-count service_id # 重置指定服务的失效计数
python consul_manager.py --reset-all-failed-count        # 重置所有服务的失效计数

注意事项:
- cleanup-duplicates: 安全清理重复实例，保留每个服务的一个健康实例
- force-cleanup: 危险操作，会清理指定服务的所有实例，请谨慎使用
- 建议定期运行 --cleanup-duplicates 来维护服务注册表的清洁
- 后台清理任务会每2分钟检查一次，失效服务需连续2次检查都是critical状态才会被清理
"""

import argparse
import sys
import asyncio
from collections import defaultdict
from utils.tll_consul import ConsulClient
from loguru import logger

def setup_logger():
    """设置日志"""
    logger.remove()
    logger.add(sys.stdout, level="INFO", format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}")

def list_all_services():
    """列出所有服务"""
    client = ConsulClient()
    services = client.get_all_services()
    
    print(f"\n=== 所有注册的服务 ({len(services)} 个) ===")
    for service_id, service_info in services.items():
        print(f"ID: {service_id}")
        print(f"  服务名: {service_info['Service']}")
        print(f"  地址: {service_info['Address']}:{service_info['Port']}")
        print(f"  标签: {service_info.get('Tags', [])}")
        print()

def list_critical_services():
    """列出所有critical状态的服务"""
    client = ConsulClient()
    critical_services = client.get_critical_services()
    
    print(f"\n=== Critical状态的服务 ({len(critical_services)} 个) ===")
    if not critical_services:
        print("没有发现critical状态的服务")
        return
        
    for service in critical_services:
        print(f"服务名: {service.get('service_name', 'N/A')}")
        print(f"  服务ID: {service.get('service_id', 'N/A')}")
        print(f"  地址: {service.get('address', 'N/A')}:{service.get('port', 'N/A')}")
        print(f"  检查ID: {service.get('check_id', 'N/A')}")
        print(f"  状态: {service.get('status', 'N/A')}")
        print(f"  输出: {service.get('output', 'N/A')}")
        print()

def cleanup_critical_services():
    """清理所有critical状态的服务"""
    client = ConsulClient()
    
    print("\n=== 开始清理critical状态的服务 ===")
    cleaned_count = client.cleanup_critical_services()
    print(f"清理完成，共清理了 {cleaned_count} 个critical状态的服务实例")

def cleanup_service_critical(service_name: str):
    """清理指定服务的critical实例"""
    client = ConsulClient()
    
    print(f"\n=== 开始清理服务 {service_name} 的critical实例 ===")
    cleaned_count = client.cleanup_critical_services(service_name=service_name)
    print(f"清理完成，共清理了 {cleaned_count} 个 {service_name} 服务的critical实例")

def cleanup_duplicate_services():
    """清理重复的服务实例，每个服务只保留一个实例"""
    client = ConsulClient()
    
    print("\n=== Consul 重复服务实例清理工具 ===")
    print("此工具将保留每个服务的一个实例，清理其他重复的实例")
    
    confirm = input("\n是否继续清理重复实例？(输入 'yes' 确认): ")
    if confirm.lower() != 'yes':
        print("操作已取消")
        return
    
    # 获取所有服务
    services = client.get_all_services()
    if not services:
        print("没有发现注册的服务")
        return
    
    # 按服务名和地址分组
    service_groups = defaultdict(list)
    
    for service_id, service_info in services.items():
        service_name = service_info['Service']
        address = f"{service_info['Address']}:{service_info['Port']}"
        
        service_groups[(service_name, address)].append({
            'id': service_id,
            'name': service_name,
            'address': address,
            'info': service_info
        })
    
    print(f"\n=== 发现的服务分组 ===")
    cleaned_count = 0
    
    for (service_name, address), instances in service_groups.items():
        print(f"\n服务: {service_name} @ {address}")
        print(f"  实例数量: {len(instances)}")
        
        if len(instances) > 1:
            # 保留第一个实例，清理其他的
            keep_instance = instances[0]
            duplicate_instances = instances[1:]
            
            print(f"  保留实例: {keep_instance['id']}")
            
            for dup in duplicate_instances:
                try:
                    client.client.agent.service.deregister(dup['id'])
                    print(f"  ✅ 已清理重复实例: {dup['id']}")
                    cleaned_count += 1
                except Exception as e:
                    print(f"  ❌ 清理实例失败 {dup['id']}: {e}")
        else:
            print(f"  ✅ 无重复实例")
    
    print(f"\n=== 清理完成 ===")
    print(f"总共清理了 {cleaned_count} 个重复的服务实例")

def force_cleanup_service(service_name: str):
    """强制清理指定服务的所有实例"""
    print(f"\n⚠️  警告：即将强制清理服务 {service_name} 的所有实例！")
    confirm = input("这是一个危险操作，是否继续？(输入 'yes' 确认): ")
    
    if confirm.lower() != 'yes':
        print("操作已取消")
        return
        
    client = ConsulClient()
    cleaned_count = client.force_cleanup_service_instances(service_name)
    print(f"强制清理完成，共清理了 {cleaned_count} 个 {service_name} 服务实例")

async def test_service_discovery():
    """测试服务发现功能"""
    client = ConsulClient()
    
    print("\n=== 测试服务发现 ===")
    print("获取健康的user_service实例...")
    addresses = await client.fetch_user_service_addresses(only_healthy=True)
    print(f"发现 {len(addresses)} 个健康的user_service实例:")
    for addr in addresses:
        print(f"  - {addr}")
    
    print("\n获取一个负载均衡的地址...")
    host, port = client.get_one_user_service_address()
    if host and port:
        print(f"负载均衡选择的地址: {host}:{port}")
    else:
        print("没有可用的服务地址")

def show_failed_count_status():
    """显示当前失效服务计数状态"""
    from utils.consul_cleanup import cleanup_task
    
    print("\n=== 失效服务计数状态 ===")
    failed_services = cleanup_task.get_failed_services_status()
    
    if not failed_services:
        print("当前没有跟踪的失效服务")
        return
    
    print(f"当前跟踪 {len(failed_services)} 个失效服务:")
    for service_id, count in failed_services.items():
        print(f"  - {service_id}: 失效{count}次")
    
    print(f"\n注意: 失效次数达到 {cleanup_task.fail_count_threshold} 次时将自动清理")

def reset_service_failed_count(service_id: str):
    """重置指定服务的失效计数"""
    from utils.consul_cleanup import cleanup_task
    
    success = cleanup_task.reset_service_fail_count(service_id)
    if success:
        print(f"✅ 已重置服务 {service_id} 的失效计数")
    else:
        print(f"❌ 服务 {service_id} 不在失效跟踪列表中")

def reset_all_failed_counts():
    """重置所有服务的失效计数"""
    from utils.consul_cleanup import cleanup_task
    
    count = cleanup_task.reset_all_fail_counts()
    print(f"✅ 已重置所有 {count} 个服务的失效计数")

def main():
    parser = argparse.ArgumentParser(description='Consul服务管理工具')
    parser.add_argument('--list', action='store_true', help='列出所有服务')
    parser.add_argument('--list-critical', action='store_true', help='列出所有critical状态的服务')
    parser.add_argument('--cleanup', action='store_true', help='清理所有critical状态的服务')
    parser.add_argument('--cleanup-service', type=str, help='清理指定服务的critical实例')
    parser.add_argument('--cleanup-duplicates', action='store_true', help='清理重复的服务实例（保留每个服务的一个实例）')
    parser.add_argument('--force-cleanup', type=str, help='强制清理指定服务的所有实例（危险操作）')
    parser.add_argument('--test-discovery', action='store_true', help='测试服务发现功能')
    parser.add_argument('--show-failed-count', action='store_true', help='显示当前失效服务计数状态')
    parser.add_argument('--reset-failed-count', type=str, help='重置指定服务的失效计数')
    parser.add_argument('--reset-all-failed-count', action='store_true', help='重置所有服务的失效计数')
    
    args = parser.parse_args()
    
    if not any(vars(args).values()):
        parser.print_help()
        return
    
    setup_logger()
    
    try:
        if args.list:
            list_all_services()
        
        if args.list_critical:
            list_critical_services()
        
        if args.cleanup:
            cleanup_critical_services()
        
        if args.cleanup_service:
            cleanup_service_critical(args.cleanup_service)
        
        if args.cleanup_duplicates:
            cleanup_duplicate_services()
        
        if args.force_cleanup:
            force_cleanup_service(args.force_cleanup)
            
        if args.test_discovery:
            asyncio.run(test_service_discovery())
            
        if args.show_failed_count:
            show_failed_count_status()
            
        if args.reset_failed_count:
            reset_service_failed_count(args.reset_failed_count)
            
        if args.reset_all_failed_count:
            reset_all_failed_counts()
            
    except Exception as e:
        logger.error(f"执行操作时发生错误: {e}")
        return 1
    
    return 0

if __name__ == "__main__":
    exit(main())
